﻿// FixIAT_dlg.cpp: 实现文件
//

#include "pch.h"
#include "afxdialogex.h"
#include "FixIAT_dlg.h"


// FixIAT_dlg 对话框

IMPLEMENT_DYNAMIC(FixIAT_dlg, CDialogEx)

FixIAT_dlg::FixIAT_dlg(CWnd* pParent /*=nullptr*/)
	: CDialogEx(IDD_DIALOG3, pParent)
	, m_rebuildIatAddr(_T("421000"))
	, m_codeSection(_T("411000"))
	, m_codeSectionSize(_T("8000"))
	, m_OEP(_T("411023"))
{

}

FixIAT_dlg::~FixIAT_dlg()
{
}

void FixIAT_dlg::DoDataExchange(CDataExchange* pDX)
{
	CDialogEx::DoDataExchange(pDX);
	DDX_Text(pDX, IDC_EDIT1, m_rebuildIatAddr);
	DDX_Text(pDX, IDC_EDIT_codeSection, m_codeSection);
	DDX_Text(pDX, IDC_EDIT_codeSectionSize, m_codeSectionSize);
	DDX_Text(pDX, IDC_EDIT_OEP, m_OEP);
}


BEGIN_MESSAGE_MAP(FixIAT_dlg, CDialogEx)
	ON_WM_CLOSE()
	ON_BN_CLICKED(IDC_BUTTON_fixIAT, &FixIAT_dlg::OnBnClickedButtonfixiat)
END_MESSAGE_MAP()


// FixIAT_dlg 消息处理程序


FixIAT_dlg g_FixIatDlg;		// 对话框
FixIatStruct g_fixIatstrut;
pIatDetailInfor g_pIat;
bool isOpenStepinto = false;

// 创建非模式对话框
void create_NoModelDialog()
{
	// 检查对话框是否已经创建
	if (!g_FixIatDlg.GetSafeHwnd())
	{
		AFX_MANAGE_STATE(AfxGetStaticModuleState());
		g_FixIatDlg.Create(IDD_DIALOG3); // 创建对话框
	}

	g_FixIatDlg.ShowWindow(SW_SHOW);	// 显示对话框
	return;
}




// 注意：非模态窗口，必须调用DestroyWindow函数来销毁窗口，不然会造成错误，原因暂时未知。
void FixIAT_dlg::OnClose()
{
	// TODO: 在此添加消息处理程序代码和/或调用默认值
	DestroyWindow();
	CDialogEx::OnClose();
}


// 开始修复（按钮）
void FixIAT_dlg::OnBnClickedButtonfixiat()
{
	// TODO: 在此添加控件通知处理程序代码
	UpdateData();
	g_fixIatstrut.codeSectionAddr = strtoll(m_codeSection.GetString(), NULL, 16);
	g_fixIatstrut.codeSectionSize = strtol(m_codeSectionSize.GetString(), NULL, 16);
	g_fixIatstrut.rebuildAddr = strtoll(m_rebuildIatAddr.GetString(), NULL, 16);
	g_fixIatstrut.OEP = strtoll(m_OEP.GetString(), NULL, 16);
	
	isOpenStepinto = true;

	// 扫描代码段
	scanCodeSection();	
	return;
}


// 扫描代码段
void scanCodeSection()
{
	/*
	* 从代码段起始位置开始扫描，遇到"0xE8"类型的call指令则停下来，判断跳转地址是不是当前节区，
	* 是则保存在一个缓存中，反之，保存到另一个缓存中。
	*  
	*/

	char ch; 
	size_t constant;
	DISASM_INSTR pInstr, tmp_pInstr;
	size_t analyIatAddr = g_fixIatstrut.rebuildAddr;
	while (0 <= g_fixIatstrut.codeSectionSize)
	{
		DbgDisasmAt(g_fixIatstrut.codeSectionAddr, &pInstr);
		DbgMemRead(g_fixIatstrut.codeSectionAddr, &ch, 1);
		if (ch == (char)0xE8 ) 
		{
			constant = pInstr.arg[0].constant;
			if (constant < analyIatAddr)  
			{
				size_t tmp_disasmAddr= g_fixIatstrut.codeSectionAddr;
				// call跳转到0xE9类型的jmp，需要重新处理
				do
				{
					DbgMemRead(constant, &ch, 1);
					if (ch != (char)0xE9) {
						break;
					}
					DbgDisasmAt(constant, &tmp_pInstr);
					tmp_disasmAddr = constant;
					constant = tmp_pInstr.arg[0].constant;					
				} while (true);		

				g_fixIatstrut.callAddr.push_back(std::make_pair(tmp_disasmAddr, constant));
			} 
			else 
			{			
				SetEIP(g_fixIatstrut.codeSectionAddr);	// 设置EIP
				DbgCmdExec("stepinto");
				return;
			}
		}

		g_fixIatstrut.codeSectionSize -= pInstr.instr_size;
		g_fixIatstrut.codeSectionAddr += pInstr.instr_size;
	}

	if (g_fixIatstrut.codeSectionSize <= 0)
	{
		// 设置EIP为OEP入口
		DbgCmdExec("pause");
		SetEIP(g_fixIatstrut.OEP);

		// 新建IAT表，并建立映射
		rebuildIat();

		// 修复分支指令
		repairBranchInstruction();
		DbgCmdExec("msg \"IAT repair completed.\"");
	}

	return;
}


// 单步回调
void sub_stepIntoCallBack(CBTYPE cbType, void* callbackInfo)
{

	if (isOpenStepinto == false) {
		return;
	}

	
	DISASM_INSTR pInstr;
	int eip = GetEIP();
	DbgDisasmAt(GetEIP(), &pInstr);
	if (pInstr.type != instr_branch) {
		DbgCmdExec("stepinto");
		return;
	}

	if (Str_IsBeginWith(pInstr.instruction, "ret") != SUCCESS_STATUS) {
		DbgCmdExec("stepinto");
		return;
	}

	size_t esp = GetESP();
	size_t buffer = 0;
	DbgMemRead(esp, &buffer, sizeof(size_t));


	char text[0x60] = { 0 };
	DbgGetModuleAt(buffer, text);
	if (text[0] == 0 || strcmp(g_fixIatstrut.mainModuleName, text) == 0) {
		DbgCmdExec("stepinto");
		return;
	}


	pIatDetailInfor pIatDI = new IatDetailInfor();
	pIatDI->callAddr = g_fixIatstrut.codeSectionAddr;
	pIatDI->trueAddr = g_fixIatstrut.codeSectionAddr;
	pIatDI->dllFuncAddr = buffer;// 把dll函数地址存放到dllFuncAddr字段中

	
	size_t getAddr = 0;
	size_t nextAddr = pIatDI->trueAddr + 5;
	size_t readAddr = sizeof(size_t) + GetESP();
	DbgMemRead(readAddr, &getAddr, sizeof(size_t));

	if (pInstr.arg[0].constant == 0)
	{
		if (nextAddr == getAddr)
		{
			pIatDI->trueAddr -= 1;
		}
		else
		{
			g_fixIatstrut.codeSectionAddr++;
			g_fixIatstrut.codeSectionSize--;
		}
	}
	else
	{
		pIatDI->isJmpInst = true;		
		static char tmpJp[] = { 0,0,0};
		DbgMemRead(nextAddr, tmpJp,sizeof(tmpJp));
		if (tmpJp[1] == (char)0xE8)
		{
			g_fixIatstrut.codeSectionAddr++;
			g_fixIatstrut.codeSectionSize--;
		}
		else if (tmpJp[2] == (char)0xE8)
		{
			g_fixIatstrut.codeSectionAddr+=2;
			g_fixIatstrut.codeSectionSize -= 2;
		}
	}
	

	g_fixIatstrut.codeSectionAddr += 5;
	g_fixIatstrut.codeSectionSize -= 5;
	g_fixIatstrut.modCallAddr.push_back(pIatDI);
		
	
	static int calc = 1;
	dprintf("===CALL===：%d,0x%x\n", calc++, pIatDI->callAddr);

	if (0 < g_fixIatstrut.codeSectionSize)
	{
		// 继续扫描代码段
		scanCodeSection();
	}
	
	return;
}





// 新建IAT表，并建立映射
void rebuildIat()
{

#define INSTUCTIONNUM	6	
	char branch[INSTUCTIONNUM] = { 0xff,0x15,0,0,0,0 };

	int totalSize = g_fixIatstrut.modCallAddr.size();
	size_t jmpTable = g_fixIatstrut.rebuildAddr;					// 新建跳转表地址
	size_t rebuildAddr = jmpTable + totalSize * INSTUCTIONNUM * 4;	// 新建IAT地址
	for (auto iatDat : g_fixIatstrut.modCallAddr)
	{

#ifdef _WIN64
		// 处理64位
		// 目标地址 = 当前地址+偏移+5
		//size_t offset = iatDat->trueAddr - val.first - 5;
		//DbgMemWrite(val.first + 1, &offset, 4);

#else
		// 处理32位						
		if (iatDat->isJmpInst == true)
		{
			branch[1] = 0x25;	// 0x25是jmp指令	
			iatDat->trueAddr = jmpTable;
			jmpTable += INSTUCTIONNUM;
			
		}  else  {
			branch[1] = 0x15;	// 0x15是call指令	
		}

		*(size_t*)(branch + 2) = rebuildAddr;
		DbgMemWrite(iatDat->trueAddr, branch, INSTUCTIONNUM);
		DbgMemWrite(rebuildAddr, &iatDat->dllFuncAddr, sizeof(size_t));

#endif
		rebuildAddr += sizeof(size_t);
	}
}


// 修复分支指令
void repairBranchInstruction()
{
	// 此函数功能：判断要修复的分支指令是call还是jmp
	for (auto val : g_fixIatstrut.callAddr)
	{
		for (auto iatDat : g_fixIatstrut.modCallAddr)
		{

			if (val.second == iatDat->callAddr || val.second == (iatDat->callAddr - 1))
			{
				iatDat->isJmpInst = true;
				// 目标地址 = 当前地址 + 偏移 + 5
				// 偏移 = 目标地址 - 当前地址 - 5
				size_t offset = iatDat->trueAddr - val.first - 5;
				DbgMemWrite(val.first + 1, &offset, 4);
			}
		}
	}

}


// 判断str的第一个子串与cutStr是否相等
bool Str_IsBeginWith(char* str, const char* cutStr)
{
	char* token ;
	strtok_s(str, " ", &token); // 获取第一个子字符串
	if (_stricmp(str, cutStr) == 0) {
		return SUCCESS_STATUS; // 成功返回0
	}
	return FAILURE_STATUS;	// 失败返回-1
}



//停止调试回调
void sub_stopDebugCallBack(CBTYPE cbType, void* callbackInfo)
{
	g_fixIatstrut.callAddr.clear();
	g_fixIatstrut.modCallAddr.clear();

	isOpenStepinto = false;
	return;
}